#! /usr/bin/env python
# -*- coding: utf-8 -*-

import os
import glob
from shutil import move
from os import listdir, rmdir

import torch
import torchvision
import torchvision.transforms as transforms
from torch.utils.data import Dataset, DataLoader
import pandas as pd
import numpy as np


class FurnaceDataset(Dataset):

    def __init__(self, path, flag='train'):
        assert flag in ['train', 'test', 'valid']
        self.flag = flag
        self.data, self.targets = self.__load_data__(path)

    def __getitem__(self, index):
        return self.data[index], self.targets[index]
        # return self.preprocess(self.data[index]), self.preprocess(self.targets[index])

    def __len__(self):
        return len(self.data)

    def __load_data__(self, path):
        if self.flag == 'train':
            df = pd.read_csv(path+'/data.csv')
        else:
            df = pd.read_csv(path+'/test_data.csv')
        x = torch.tensor(df.iloc[:, 7:15].to_numpy(dtype=np.float32))
        y = torch.tensor(df.iloc[:, 1:7].to_numpy(dtype=np.float32))
        # print(torch.mean(x), torch.std(x))
        # print(torch.mean(y), torch.std(y))
        # x = (x - torch.mean(x)) / torch.std(x)
        # y = (y - torch.mean(y)) / torch.std(y)
        return x, y

    def preprocess(self, data):
        return data


# Based on https://github.com/alecwangcq/GraSP/blob/master/utils/data_utils.py
def get_transforms(config):
    dataset = config.dataset
    network = config.network
    transform_train = None
    transform_test = None
    if dataset == 'mnist':
        # transforms.Normalize((0.1307,), (0.3081,))
        t = transforms.Normalize((0.5,), (0.5,))
        transform_train = transforms.Compose([transforms.ToTensor(), t])
        transform_test = transforms.Compose([transforms.ToTensor(), t])

    assert transform_test is not None and transform_train is not None, 'Error, no dataset %s' % dataset
    return transform_train, transform_test


def get_dataloader(config, num_workers=0, trainset_shuffle=True):
    dataset = config.dataset
    train_batch_size = config.batch_size
    test_batch_size = 256
    root = config.dp

    # transform_train, transform_test = get_transforms(config)
    trainset, testset = None, None
    # if dataset == 'mnist':
    #     trainset = torchvision.datasets.MNIST(root=root + '/mnist', train=True, download=True, transform=transform_train)
    #     testset = torchvision.datasets.MNIST(root=root + '/mnist', train=False, download=True, transform=transform_test)
    if dataset == 'furnace86':
        trainset = FurnaceDataset(root+'/furnace86/csv', 'train')
        testset = FurnaceDataset(root+'/furnace86/csv', 'test')
        print(len(trainset))
        print((trainset[0]))
        print(len(testset))
        print((testset[0]))

    assert trainset is not None and testset is not None, 'Error, no dataset %s' % dataset
    trainloader = torch.utils.data.DataLoader(trainset, batch_size=train_batch_size, shuffle=trainset_shuffle, num_workers=num_workers)
    testloader = torch.utils.data.DataLoader(testset, batch_size=test_batch_size, shuffle=False, num_workers=num_workers)

    return trainloader, testloader
